package org.cybzacg.encryption.impl;

import org.cybzacg.encryption.core.AbstractEncryptor;
import org.cybzacg.encryption.context.EncryptionContext;
import org.cybzacg.encryption.model.ValidationResult;
import org.cybzacg.encryption.enums.EncryptionAlgorithm;
import org.cybzacg.encryption.exception.EncryptionException;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.SecureRandom;

/**
 * AES加密器实现
 * 支持AES-128、AES-192、AES-256加密算法
 * 
 * @author encryption team
 * @version 1.0
 */
public class AESEncryptor extends AbstractEncryptor {

    private static final String ALGORITHM_NAME = "AES";
    private static final String TRANSFORMATION = "AES/CBC/PKCS5Padding";
    private static final int IV_LENGTH = 16; // AES块大小

    private volatile SecretKey secretKey;
    private volatile IvParameterSpec ivSpec;

    /**
     * 构造函数
     */
    public AESEncryptor() {
        super(EncryptionAlgorithm.AES);
    }

    /**
     * 构造函数（指定密钥长度）
     * 
     * @param keySize 密钥长度（128、192、256）
     */
    public AESEncryptor(int keySize) {
        super(EncryptionAlgorithm.AES);
        metadata.put("keySize", keySize);
    }

    @Override
    protected void doInitialize(EncryptionContext context) {
        try {
            // 获取或生成密钥
            byte[] keyBytes = context.getKey();
            if (keyBytes == null) {
                // 生成新密钥
                KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM_NAME);
                keyGen.init(getKeySize(context));
                secretKey = keyGen.generateKey();
                keyBytes = secretKey.getEncoded();
            } else {
                // 使用提供的密钥
                if (!algorithm.isValidKeySize(keyBytes.length * 8)) {
                    throw EncryptionException.invalidKey(
                            "Invalid AES key size: " + (keyBytes.length * 8) + " bits",
                            context);
                }
                secretKey = new SecretKeySpec(keyBytes, ALGORITHM_NAME);
            }

            // 获取或生成IV
            byte[] ivBytes = context.getIv();
            if (ivBytes == null) {
                // 生成随机IV
                byte[] iv = new byte[IV_LENGTH];
                new SecureRandom().nextBytes(iv);
                ivSpec = new IvParameterSpec(iv);
            } else {
                // 使用提供的IV
                if (ivBytes.length != IV_LENGTH) {
                    throw EncryptionException.invalidDataFormat(
                            "Invalid IV length: " + ivBytes.length + " (expected: " + IV_LENGTH + ")",
                            context);
                }
                ivSpec = new IvParameterSpec(ivBytes);
            }

            // 更新元数据
            metadata.put("transformation", TRANSFORMATION);
            metadata.put("ivLength", IV_LENGTH);
            metadata.put("initialized", true);

        } catch (Exception e) {
            throw new EncryptionException("Failed to initialize AES encryptor", e);
        }
    }

    @Override
    protected byte[] doEncrypt(byte[] data, EncryptionContext context) {
        try {
            // 每次加密都创建新的 Cipher 实例（保证线程安全）
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivSpec);
            return cipher.doFinal(data);
        } catch (Exception e) {
            throw new EncryptionException("AES encryption failed", e);
        }
    }

    @Override
    protected byte[] doDecrypt(byte[] encryptedData, EncryptionContext context) {
        try {
            // 每次解密都创建新的 Cipher 实例（保证线程安全）
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            cipher.init(Cipher.DECRYPT_MODE, secretKey, ivSpec);
            return cipher.doFinal(encryptedData);
        } catch (Exception e) {
            throw new EncryptionException("AES decryption failed", e);
        }
    }

    @Override
    protected ValidationResult doValidate(EncryptionContext context) {
        try {
            // 验证密钥
            byte[] keyBytes = context.getKey();
            if (keyBytes != null) {
                int keySizeBits = keyBytes.length * 8;
                if (!algorithm.isValidKeySize(keySizeBits)) {
                    return ValidationResult.failure(
                            "Invalid AES key size",
                            "Key size: " + keySizeBits + " bits (supported: 128, 192, 256)",
                            context);
                }
            }

            // 验证IV
            byte[] ivBytes = context.getIv();
            if (ivBytes != null && ivBytes.length != IV_LENGTH) {
                return ValidationResult.failure(
                        "Invalid IV length",
                        "IV length: " + ivBytes.length + " (expected: " + IV_LENGTH + ")",
                        context);
            }

            return ValidationResult.success("AES validation passed", context);

        } catch (Exception e) {
            return ValidationResult.failure("AES validation error", e.getMessage(), context);
        }
    }

    @Override
    protected boolean doSupports(EncryptionContext context) {
        if (context == null) {
            return false;
        }

        // 检查算法匹配
        if (context.getAlgorithm() != EncryptionAlgorithm.AES) {
            return false;
        }

        // 检查加密类型匹配
        if (context.getEncryptionType() != algorithm.getType()) {
            return false;
        }

        // AES是对称加密，需要密钥
        if (context.getKey() == null) {
            return false;
        }

        return true;
    }

    @Override
    protected void doReset() {
        // Cipher 每次使用时创建新实例，无需重置
        // 只需确保密钥和 IV 仍然有效
    }

    @Override
    protected void doClose() {
        secretKey = null;
        ivSpec = null;
        metadata.put("initialized", false);
    }

    /**
     * 获取密钥长度
     * 
     * @param context 上下文
     * @return 密钥长度
     */
    private int getKeySize(EncryptionContext context) {
        // 优先使用上下文中的密钥长度
        Integer keySize = (Integer) context.getAttribute("keySize");
        if (keySize != null) {
            return keySize;
        }

        // 使用元数据中的密钥长度
        keySize = (Integer) metadata.get("keySize");
        if (keySize != null) {
            return keySize;
        }

        // 使用配置中的密钥长度
        if (context.getConfig() != null) {
            return context.getConfig().getKeySize();
        }

        // 默认使用256位
        return 256;
    }

    /**
     * 生成随机密钥
     * 
     * @param keySize 密钥长度
     * @return 密钥字节数组
     */
    public static byte[] generateKey(int keySize) {
        try {
            KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM_NAME);
            keyGen.init(keySize);
            SecretKey secretKey = keyGen.generateKey();
            return secretKey.getEncoded();
        } catch (Exception e) {
            throw new EncryptionException("Failed to generate AES key", e);
        }
    }

    /**
     * 生成随机IV
     * 
     * @return IV字节数组
     */
    public static byte[] generateIV() {
        byte[] iv = new byte[IV_LENGTH];
        new SecureRandom().nextBytes(iv);
        return iv;
    }

    /**
     * 获取当前密钥
     * 
     * @return 密钥字节数组，如果未初始化返回null
     */
    public byte[] getCurrentKey() {
        return secretKey != null ? secretKey.getEncoded() : null;
    }

    /**
     * 获取当前IV
     * 
     * @return IV字节数组，如果未初始化返回null
     */
    public byte[] getCurrentIV() {
        return ivSpec != null ? ivSpec.getIV() : null;
    }

    @Override
    public String getDescription() {
        return "AES (Advanced Encryption Standard) symmetric encryption algorithm";
    }
}
